#include "rbtree.h"

template<class K, class V>
void RedBlackTree<K,V>::leftRotate(std::shared_ptr<Node<K,V>> x) {
	auto y = x->right;
	x->right = y->left;
	if (y->left != nullptr)
		y->left->parent = x;
	y->parent = x->parent;
	if (!x->parent.lock())
		root = y;
	else if (x == x->parent.lock()->left)
		x->parent.lock()->left = y;
	else
		x->parent.lock()->right = y;
	y->left = x;
	x->parent = y;
}

template<class K, class V>
void RedBlackTree<K,V>::rightRotate(std::shared_ptr<Node<K,V>> y) {
	auto x = y->left;
	y->left = x->right;
	if (x->right != nullptr)
		x->right->parent = y;
	x->parent = y->parent;
	if (y->parent.lock() == nullptr)
		root = x;
	else if (y == y->parent.lock()->right)
		y->parent.lock()->right = x;
	else
		y->parent.lock()->left = x;
	x->right = y;
	y->parent = x;
}

template<class K, class V>
void RedBlackTree<K,V>::insertFixup(std::shared_ptr<Node<K,V>> z) {
	while (z->parent.lock() != nullptr && z->parent.lock()->color == RED) {
		if (z->parent.lock() == z->parent.lock()->parent.lock()->left) {
			auto y = z->parent.lock()->parent.lock()->right;
			if (y != nullptr && y->color == RED) {
				z->parent.lock()->color = BLACK;
				y->color = BLACK;
				z->parent.lock()->parent.lock()->color = RED;
				z = z->parent.lock()->parent.lock();
			} else {
				if (z == z->parent.lock()->right) {
					z = z->parent.lock();
					leftRotate(z);
				}
				z->parent.lock()->color = BLACK;
				z->parent.lock()->parent.lock()->color = RED;
				rightRotate(z->parent.lock()->parent.lock());
			}
		} else {
			auto y = z->parent.lock()->parent.lock()->left;
			if (y != nullptr && y->color == RED) {
				z->parent.lock()->color = BLACK;
				y->color = BLACK;
				z->parent.lock()->parent.lock()->color = RED;
				z = z->parent.lock()->parent.lock();
			} else {
				if (z == z->parent.lock()->left) {
					z = z->parent.lock();
					rightRotate(z);
				}
				z->parent.lock()->color = BLACK;
				z->parent.lock()->parent.lock()->color = RED;
				leftRotate(z->parent.lock()->parent.lock());
			}
		}
	}
	root->color = BLACK;
}

template<class K, class V>
void RedBlackTree<K,V>::insert(K key, V value) {
	auto z = std::make_shared<Node<K,V>>(std::make_pair(key, value));
	std::shared_ptr<Node<K,V>> y = nullptr;
	std::shared_ptr<Node<K,V>> x = root;
	
	while (x != nullptr) {
		y = x;
		if (z->data.first < x->data.first)
			x = x->left;
		else
			x = x->right;
	}
	
	z->parent = y;
	if (y == nullptr)
		root = z;
	else if (z->data.first < y->data.first)
		y->left = z;
	else
		y->right = z;
	
	insertFixup(z);
}

template<class K, class V>
void RedBlackTree<K,V>::inorderHelper(std::shared_ptr<Node<K,V>> node) {
	if (node != nullptr) {
		inorderHelper(node->left);
		std::cout << "(" << node->data.first << ", " << node->data.second << ") ";
		inorderHelper(node->right);
	}
}

template<class K, class V>
void RedBlackTree<K,V>::inorder() {
	inorderHelper(root);
}

template<typename K, typename V>
void RedBlackTree<K,V>::deleteFixup(std::shared_ptr<Node<K,V>> x) {
	while (x != root && (x == nullptr || x->color == BLACK)) {
		if (x == x->parent.lock()->left) {
			auto w = x->parent.lock()->right;
			if (w->color == RED) {
				w->color = BLACK;
				x->parent.lock()->color = RED;
				leftRotate(x->parent.lock());
				w = x->parent.lock()->right;
			}
			if ((w->left == nullptr || w->left->color == BLACK) &&
				(w->right == nullptr || w->right->color == BLACK)) {
				w->color = RED;
				x = x->parent.lock();
			} else {
				if (w->right == nullptr || w->right->color == BLACK) {
					if (w->left != nullptr)
						w->left->color = BLACK;
					w->color = RED;
					rightRotate(w);
					w = x->parent.lock()->right;
				}
				w->color = x->parent.lock()->color;
				x->parent.lock()->color = BLACK;
				if (w->right != nullptr)
					w->right->color = BLACK;
				leftRotate(x->parent.lock());
				x = root;
			}
		} else {
			auto w = x->parent.lock()->left;
			if (w->color == RED) {
				w->color = BLACK;
				x->parent.lock()->color = RED;
				rightRotate(x->parent.lock());
				w = x->parent.lock()->left;
			}
			if ((w->right == nullptr || w->right->color == BLACK) &&
				(w->left == nullptr || w->left->color == BLACK)) {
				w->color = RED;
				x = x->parent.lock();
			} else {
				if (w->left == nullptr || w->left->color == BLACK) {
					if (w->right != nullptr)
						w->right->color = BLACK;
					w->color = RED;
					leftRotate(w);
					w = x->parent.lock()->left;
				}
				w->color = x->parent.lock()->color;
				x->parent.lock()->color = BLACK;
				if (w->left != nullptr)
					w->left->color = BLACK;
				rightRotate(x->parent.lock());
				x = root;
			}
		}
	}
	if (x != nullptr)
		x->color = BLACK;
}

template<typename K, typename V>
void RedBlackTree<K,V>::transplant(std::shared_ptr<Node<K,V>> u, std::shared_ptr<Node<K,V>> v) {
	if (u->parent.lock() == nullptr)
		root = v;
	else if (u == u->parent.lock()->left)
		u->parent.lock()->left = v;
	else
		u->parent.lock()->right = v;
	if (v != nullptr)
		v->parent = u->parent;
}

template<typename K, typename V>
std::shared_ptr<Node<K,V>> RedBlackTree<K,V>::minimum(std::shared_ptr<Node<K,V>> x) {
	while (x->left != nullptr)
		x = x->left;
	return x;
}

template<typename K, typename V>
void RedBlackTree<K,V>::remove(K key) {
	auto z = root;
	while (z != nullptr) {
		if (key == z->data.first)
			break;
		else if (key < z->data.first)
			z = z->left;
		else
			z = z->right;
	}
	if (z == nullptr)
		return;
	
	auto y = z;
	std::shared_ptr<Node<K,V>> x;
	Color y_original_color = y->color;
	if (z->left == nullptr) {
		x = z->right;
		transplant(z, z->right);
	} else if (z->right == nullptr) {
		x = z->left;
		transplant(z, z->left);
	} else {
		y = minimum(z->right);
		y_original_color = y->color;
		x = y->right;
		if (y->parent.lock() == z)
			x->parent = y;
		else {
			transplant(y, y->right);
			y->right = z->right;
			y->right->parent = y;
		}
		transplant(z, y);
		y->left = z->left;
		y->left->parent = y;
		y->color = z->color;
	}
	if (y_original_color == BLACK)
		deleteFixup(x);
}

template<typename K, typename V>
void RedBlackTree<K,V>::modify(K key, V new_value) {
	auto node = root;
	while (node != nullptr) {
		if (key == node->data.first) {
			node->data.second = new_value;
			return;
		} else if (key < node->data.first)
			node = node->left;
		else
			node = node->right;
	}
}




//int main() {
//	RedBlackTree<std::string, int> rbt;
//	rbt.insert("7", 7);
//	rbt.insert("3", 3);
//	rbt.insert("18", 18);
//	rbt.insert("10", 10);
//	rbt.insert("22", 22);
//	rbt.insert("8", 8);
//	rbt.insert("11", 11);
//	
//	std::cout << "中序遍历结果: ";
//	rbt.inorder();
//	std::cout << std::endl;
//	rbt.modify("8", 9);
//	std::cout << "修改后中序遍历结果: ";
//	rbt.inorder();
//	std::cout << std::endl;
//	rbt.remove("8");
//	std::cout << "删除8结点后中序遍历结果: ";
//	rbt.inorder();
//	std::cout << std::endl;
//	for (auto it = rbt.begin(); it != rbt.end(); ++it) {
//		std::cout << "(" << (*it).first << ", " << (*it).second << ") ";
//	}
//	return 0;
//}
